#!/usr/bin/env python3
"""
统一交易系统 - 支持监控、实盘、模拟、回测
"""
import sys
import os
import time
from typing import Dict, List

# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from strategy_executor import StrategyManager, ExecutionMode

class TradingSystem:
    """统一交易系统"""
    
    def __init__(self):
        self.manager = StrategyManager()
        self.running = True
    
    @staticmethod
    def normalize_stock_code(code: str) -> str:
        """
        标准化股票代码，自动添加交易所后缀
        
        规则:
        上海交易所 (.SH):
        - 60xxxx: 主板股票
        - 68xxxx: 主板B股
        - 688xxx: 科创板
        - 50xxxx: 基金
        - 51xxxx: 上海ETF
        
        深圳交易所 (.SZ):
        - 00xxxx: 主板股票
        - 30xxxx: 创业板
        - 200xxx: B股
        - 15xxxx: 深圳ETF
        - 159xxx: 深圳ETF
        
        Args:
            code: 股票代码，可以是6位数字或带后缀（.SH/.SZ/.sh/.sz）
            
        Returns:
            标准化后的股票代码（大写后缀）
        """
        code = code.strip().upper()
        
        # 如果已经包含后缀，统一转大写并返回
        if '.SH' in code or '.SZ' in code:
            return code
        
        # 只保留数字
        code_num = ''.join(filter(str.isdigit, code))
        
        if len(code_num) != 6:
            # 如果不是6位数字，返回原始输入（已转大写）
            return code
        
        # 根据前缀判断交易所
        prefix = code_num[:2]
        prefix3 = code_num[:3]
        
        # 上海交易所
        if prefix in ['60', '68', '50', '51'] or prefix3 == '688':
            return f"{code_num}.SH"
        # 深圳交易所（包括ETF）
        elif prefix in ['00', '30', '20', '15'] or prefix3 == '159':
            return f"{code_num}.SZ"
        else:
            # 无法识别，返回原始输入
            return code
        
    def show_main_menu(self):
        """显示主菜单"""
        print("\n" + "="*60)
        print("🚀 统一交易系统")
        print("="*60)
        print("1. 策略监控模式")
        print("2. 模拟交易模式") 
        print("3. 实盘交易模式")
        print("4. 历史回测模式")
        print("5. 管理运行中的策略")
        print("6. 查看可用策略")
        print("0. 退出系统")
        print("="*60)
    
    def show_strategy_selection(self, allow_multiple=False):
        """显示策略选择菜单"""
        strategies = self.manager.get_available_strategies()
        if not strategies:
            print("❌ 未找到可用策略")
            return None
        
        print("\n📋 可用策略列表:")
        for i, strategy in enumerate(strategies, 1):
            print(f"{i}. {strategy}")
        
        if allow_multiple:
            print("\n💡 提示: 可以选择多个策略，用逗号分隔 (如: 1,3,5)")
        
        while True:
            try:
                if allow_multiple:
                    choice = input(f"\n请选择策略 (1-{len(strategies)}, 支持多选): ").strip()
                    if ',' in choice:
                        # 多选
                        indices = [int(x.strip()) - 1 for x in choice.split(',')]
                        selected = [strategies[i] for i in indices if 0 <= i < len(strategies)]
                        if selected:
                            return selected
                        print("❌ 无效选择，请重新输入")
                    else:
                        # 单选
                        if choice.isdigit():
                            idx = int(choice) - 1
                            if 0 <= idx < len(strategies):
                                return [strategies[idx]]
                else:
                    choice = input(f"\n请选择策略 (1-{len(strategies)}): ").strip()
                    if choice.isdigit():
                        idx = int(choice) - 1
                        if 0 <= idx < len(strategies):
                            return strategies[idx]
                print("❌ 无效选择，请重新输入")
            except KeyboardInterrupt:
                return None
            except ValueError:
                print("❌ 输入格式错误，请重新输入")
    
    def start_monitor_mode(self):
        """启动监控模式"""
        print("\n🔍 策略监控模式")
        print("监控模式：只输出信号，不执行交易")
        
        strategies = self.show_strategy_selection(allow_multiple=True)
        if not strategies:
            return
        
        if not isinstance(strategies, list):
            strategies = [strategies]
        
        # 获取股票和板块输入（允许为空，只监控持仓）
        target_stocks, target_sectors = self.get_stocks_and_sectors_input()
        
        try:
            executors = []
            for strategy in strategies:
                executor_id, executor = self.manager.create_executor(
                    strategy, ExecutionMode.MONITOR,
                    target_stocks=target_stocks,
                    target_sectors=target_sectors
                )
                executor.start()
                executors.append((executor_id, executor))
            
            print(f"\n✅ 启动监控模式 - 策略: {', '.join(strategies)}")
            if target_stocks:
                print(f"   监控股票: {', '.join(target_stocks)}")
            if target_sectors:
                print(f"   监控板块: {', '.join(target_sectors)}")
            print("按 Ctrl+C 停止监控")
            
            # 等待用户中断
            try:
                while True:
                    time.sleep(1)
            except KeyboardInterrupt:
                print("\n🛑 停止监控模式")
                for executor_id, executor in executors:
                    executor.stop()
                
        except Exception as e:
            print(f"❌ 启动监控模式失败: {str(e)}")
    
    def get_stocks_and_sectors_input(self):
        """获取用户输入的股票代码和板块"""
        print("\n📊 请配置监控范围")
        print("=" * 50)
        print("💡 提示: 系统会自动监控已持仓的股票")
        print("        如果不输入任何股票/板块，将只监控已持仓")
        
        # 获取股票代码
        print("\n1️⃣  输入股票代码（可输入多个，用逗号或空格分隔）")
        print("   示例: 002594 或 002594,600519 或 002594 600519")
        print("   直接回车跳过")
        stocks_input = input("股票代码: ").strip()
        
        target_stocks = []
        if stocks_input:
            # 支持逗号和空格分隔
            stocks_input = stocks_input.replace(',', ' ')
            raw_stocks = [s.strip() for s in stocks_input.split() if s.strip()]
            
            # 自动补全后缀
            for stock in raw_stocks:
                if '.' not in stock:
                    # 根据代码开头判断市场
                    if stock.startswith('6') or stock.startswith('5') or stock.startswith('9'):
                        stock = stock + '.SH'
                    else:
                        stock = stock + '.SZ'
                target_stocks.append(stock.upper())
            
            print(f"   ✓ 已添加 {len(target_stocks)} 只股票: {', '.join(target_stocks)}")
        
        # 获取板块
        print("\n2️⃣  输入板块名称（可输入多个，用逗号或空格分隔）")
        print("   常用板块: 上证, 科创, 创业板, 沪深300")
        print("   ⚠️  注意: 板块股票数量较多（如上证2000+只），下载数据耗时较长")
        print("   直接回车跳过")
        sectors_input = input("板块名称: ").strip()
        
        target_sectors = []
        if sectors_input:
            # 支持逗号和空格分隔
            sectors_input = sectors_input.replace(',', ' ')
            target_sectors = [s.strip() for s in sectors_input.split() if s.strip()]
            print(f"   ✓ 已添加 {len(target_sectors)} 个板块: {', '.join(target_sectors)}")
        
        # 允许全部跳过（只监控持仓）
        if not target_stocks and not target_sectors:
            print("\n💡 将只监控已持仓的股票")
        
        print("\n" + "=" * 50)
        return target_stocks, target_sectors
    
    def start_simulate_mode(self):
        """启动模拟交易模式"""
        print("\n💰 模拟交易模式")
        print("模拟交易：使用虚拟资金进行交易")
        
        strategy = self.show_strategy_selection()
        if not strategy:
            return
        
        # 获取股票和板块输入（允许为空，只监控持仓）
        target_stocks, target_sectors = self.get_stocks_and_sectors_input()
        
        # 获取初始资金
        try:
            capital = float(input("\n请输入初始资金 (默认100000): ") or "100000")
        except ValueError:
            capital = 100000
        
        try:
            executor_id, executor = self.manager.create_executor(
                strategy, ExecutionMode.SIMULATE, 
                initial_capital=capital,
                target_stocks=target_stocks,
                target_sectors=target_sectors
            )
            
            print(f"\n✅ 启动模拟交易 - 策略: {strategy}, 初始资金: {capital:,.2f}元")
            if target_stocks:
                print(f"   监控股票: {', '.join(target_stocks)}")
            if target_sectors:
                print(f"   监控板块: {', '.join(target_sectors)}")
            print("按 Ctrl+C 停止模拟交易")
            
            executor.start()
            
            try:
                while True:
                    time.sleep(1)
            except KeyboardInterrupt:
                print("\n🛑 停止模拟交易")
                executor.stop()
                
        except Exception as e:
            print(f"❌ 启动模拟交易失败: {str(e)}")
    
    def start_live_trade_mode(self):
        """启动实盘交易模式"""
        print("\n⚠️ 实盘交易模式")
        print("⚠️ 警告：此模式将执行真实交易！")
        
        confirm = input("确认启动实盘交易模式？(输入 'YES' 确认): ")
        if confirm != 'YES':
            print("❌ 已取消实盘交易")
            return
        
        strategy = self.show_strategy_selection()
        if not strategy:
            return
        
        # 获取股票和板块输入（允许为空，只监控持仓）
        target_stocks, target_sectors = self.get_stocks_and_sectors_input()
        
        try:
            executor_id, executor = self.manager.create_executor(
                strategy, ExecutionMode.LIVE_TRADE,
                target_stocks=target_stocks,
                target_sectors=target_sectors
            )
            
            print(f"\n✅ 启动实盘交易 - 策略: {strategy}")
            if target_stocks:
                print(f"   监控股票: {', '.join(target_stocks)}")
            if target_sectors:
                print(f"   监控板块: {', '.join(target_sectors)}")
            print("⚠️ 实盘交易已启动，请谨慎操作！")
            print("按 Ctrl+C 停止实盘交易")
            
            executor.start()
            
            try:
                while True:
                    time.sleep(1)
            except KeyboardInterrupt:
                print("\n🛑 停止实盘交易")
                executor.stop()
                
        except Exception as e:
            print(f"❌ 启动实盘交易失败: {str(e)}")
    
    def get_available_backtest_systems(self):
        """自动发现可用的回测系统"""
        import os
        import importlib.util
        
        backtest_systems = {}
        backtest_dir = "backtest_systems"
        
        if not os.path.exists(backtest_dir):
            return backtest_systems
        
        for filename in os.listdir(backtest_dir):
            if filename.endswith('.py') and not filename.startswith('__'):
                # 提取文件名（去掉.py扩展名）
                system_name = filename[:-3]
                backtest_systems[system_name] = filename
        
        return backtest_systems
    
    def start_backtest_mode(self):
        """启动回测模式"""
        print("\n📊 历史回测模式")
        
        # 自动发现回测系统
        backtest_systems = self.get_available_backtest_systems()
        
        if not backtest_systems:
            print("❌ 未找到可用的回测系统")
            return
        
        # 显示回测类型选择
        print("\n📋 可用回测系统:")
        system_list = list(backtest_systems.keys())
        for i, system_name in enumerate(system_list, 1):
            print(f"{i}. {system_name}")
        
        while True:
            try:
                choice = input(f"请选择回测系统 (1-{len(system_list)}): ").strip()
                if choice.isdigit() and 1 <= int(choice) <= len(system_list):
                    selected_system = system_list[int(choice) - 1]
                    break
                print("❌ 无效选择，请重新输入")
            except KeyboardInterrupt:
                return
        
        strategy = self.show_strategy_selection()
        if not strategy:
            return
        
        # 获取回测参数
        try:
            start_date = input("请输入开始日期 (YYYYMMDD, 默认20250101): ").strip() or "20250101"
            end_date = input("请输入结束日期 (YYYYMMDD, 默认系统当前日期): ").strip()
            if not end_date:
                from datetime import datetime
                end_date = datetime.now().strftime('%Y%m%d')
            
            print("\n💡 股票代码输入说明:")
            print("   • 可以只输入6位代码: 601899")
            print("   • 可以带后缀: 601899.SH 或 601899.sh")
            print("   • 支持多个股票: 601899,000001,300750")
            print("   • 支持混合格式: 601899,000001.SZ,300750.sz")
            print("   • 支持ETF: 510300(上证50ETF), 159915(创业板ETF)")
            
            stock_input = input("\n请输入股票代码 (用逗号分隔, 默认000001): ").strip() or "000001"
            
            # 统一中英文逗号
            stock_input = stock_input.replace('，', ',')
            
            # 标准化股票代码，自动添加交易所后缀
            stock_list = [self.normalize_stock_code(s.strip()) for s in stock_input.split(',')]
            
            # 显示标准化后的股票代码
            print(f"\n✅ 标准化后的股票代码: {', '.join(stock_list)}")
        except Exception as e:
            print(f"❌ 参数输入错误: {str(e)}")
            return
        
        try:
            # 创建回测结果保存目录
            import os
            from datetime import datetime
            results_dir = "backtest_results"
            if not os.path.exists(results_dir):
                os.makedirs(results_dir)
            
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            result_prefix = f"{results_dir}/{strategy}_{timestamp}"
            
            # 根据选择的回测系统确定执行模式
            if selected_system == "分时回测":
                execution_mode = ExecutionMode.MINUTE_BACKTEST
            else:
                # 对于其他回测系统，默认使用分时回测模式
                execution_mode = ExecutionMode.MINUTE_BACKTEST
            
            executor_id, executor = self.manager.create_executor(
                strategy, execution_mode,
                strategy_file=strategy,
                stock_list=stock_list,
                start_date=start_date,
                end_date=end_date,
                result_prefix=result_prefix
            )
            
            print(f"\n✅ 开始回测 - 策略: {strategy}")
            print(f"回测系统: {selected_system}")
            print(f"回测期间: {start_date} 至 {end_date}")
            print(f"股票列表: {', '.join(stock_list)}")
            print(f"结果保存: {result_prefix}_*.png")
            
            executor.execute()  # 回测是同步执行
            
            print(f"\n🎉 回测完成！结果已保存到: {results_dir}/")
            print("生成的文件:")
            for file in os.listdir(results_dir):
                if file.startswith(f"{strategy}_{timestamp}"):
                    print(f"  - {file}")
                
        except Exception as e:
            print(f"❌ 回测失败: {str(e)}")
    
    def manage_running_strategies(self):
        """管理运行中的策略"""
        status = self.manager.get_executor_status()
        if not status:
            print("📭 当前没有运行中的策略")
            return
        
        print("\n📋 运行中的策略:")
        for i, (executor_id, info) in enumerate(status.items(), 1):
            print(f"{i}. {info['strategy']} ({info['mode']}) - {'运行中' if info['running'] else '已停止'}")
        
        while True:
            try:
                choice = input("\n请选择操作 (1-停止策略, 0-返回主菜单): ").strip()
                if choice == '0':
                    break
                elif choice.isdigit():
                    idx = int(choice) - 1
                    if 0 <= idx < len(status):
                        executor_id = list(status.keys())[idx]
                        self.manager.stop_executor(executor_id)
                        print(f"✅ 已停止策略: {executor_id}")
                        break
                print("❌ 无效选择")
            except KeyboardInterrupt:
                break
    
    def show_available_strategies(self):
        """显示可用策略"""
        strategies = self.manager.get_available_strategies()
        if not strategies:
            print("📭 未找到可用策略")
            return
        
        print("\n📋 可用策略列表:")
        for i, strategy in enumerate(strategies, 1):
            print(f"{i}. {strategy}")
        
        input("\n按回车键返回主菜单...")
    
    def run(self):
        """运行主循环"""
        print("🚀 欢迎使用统一交易系统！")
        
        while self.running:
            try:
                self.show_main_menu()
                choice = input("请选择操作 (0-6): ").strip()
                
                if choice == '0':
                    print("👋 感谢使用，再见！")
                    self.manager.stop_all_executors()
                    self.running = False
                elif choice == '1':
                    self.start_monitor_mode()
                elif choice == '2':
                    self.start_simulate_mode()
                elif choice == '3':
                    self.start_live_trade_mode()
                elif choice == '4':
                    self.start_backtest_mode()
                elif choice == '5':
                    self.manage_running_strategies()
                elif choice == '6':
                    self.show_available_strategies()
                else:
                    print("❌ 无效选择，请重新输入")
                    
            except KeyboardInterrupt:
                print("\n\n👋 系统正在退出...")
                self.manager.stop_all_executors()
                self.running = False
            except Exception as e:
                print(f"❌ 系统错误: {str(e)}")
                time.sleep(1)

def main():
    """主函数"""
    try:
        system = TradingSystem()
        system.run()
    except Exception as e:
        print(f"❌ 系统启动失败: {str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()
